#include <mpi.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <algorithm>

#define LEN_HOSTNAME 32
#define MAX_NUM_NODE 32

int main(int argc, char** argv){
    int rank, nproc;
    int TAG = 100;
    MPI_Request send_request, recv_request;
    MPI_Status send_status, recv_status;
    MPI_Init(&argc,&argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nproc);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    char hostname[LEN_HOSTNAME] = {0};
    char all_hostname[LEN_HOSTNAME*MAX_NUM_NODE] = {0};
    int node_flag[MAX_NUM_NODE] = {0};
    int group_split[MAX_NUM_NODE*2] = {0};
    MPI_Group node_groups[MAX_NUM_NODE];
    MPI_Comm node_comms[MAX_NUM_NODE];
    int head_group_incl[MAX_NUM_NODE];
    MPI_Group head_group;
    MPI_Comm head_comm;
    int node_cnt = 0;
    gethostname(hostname, sizeof(hostname));
    if (rank==0){
        printf(">>>>>>>>>>>>>>>>>>>>gather hostname...\n");
    }
    MPI_Gather(hostname, LEN_HOSTNAME, MPI_CHAR, all_hostname, LEN_HOSTNAME, MPI_CHAR, 0, MPI_COMM_WORLD);
    if (rank==0){
        printf("rank & hostname:\n");
        for (int i = 0; i < nproc; i++){
            int head = i*LEN_HOSTNAME;
            printf("rank=%d\thostname=%s\n",i,all_hostname+head);
        }
        int head = 0;
        int left_procs = nproc;
        while (left_procs > 0){
            while (node_flag[head]>0){
                head++;
            }
            for (int i = head; i < nproc; i++){
                if (strcmp(all_hostname+(head*LEN_HOSTNAME),all_hostname+(i*LEN_HOSTNAME))==0){
                    node_flag[i] = node_cnt;
                    left_procs--;
                }
            }
            node_cnt++;
        }
        // for (int i = 0; i<nproc; i++){
        //     printf("i=%d\tnode_flag[i]=%d\n",i,node_flag[i]);
        // }
        head = 0;
        for (int i = 0; i<node_cnt; i++){
            int tail = head+1;
            for (int j = 0; j<nproc; j++){
                if (node_flag[j] == i){
                    group_split[head]++;
                    group_split[tail++] = j;
                }
            }
            head = tail;
        }
        // printf("group_split:\t");
        // for (int i = 0; i<head; i++){
        //     printf("%d\t",group_split[i]);
        // }
        // printf("\n");
        head = 0;
        int current_flag = 0;
        while (group_split[head]>0){
            printf("node_group#%d:\t", current_flag);
            for (int i = head+1; i<=head+group_split[head]; i++){
                printf("%d\t",group_split[i]);
            }
            printf("\n");
            head = head+group_split[head]+1;
            current_flag++;
        }
    }
    int my_flag;
    MPI_Scatter(node_flag, 1, MPI_INT, &my_flag, 1, MPI_INT, 0, MPI_COMM_WORLD);
    // MPI_Bcast(&node_cnt, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(group_split,MAX_NUM_NODE*2,MPI_INT,0,MPI_COMM_WORLD);
    // printf("rank=%d\tmy_flag=%d\n",rank,my_flag);
    if (rank==0){
        printf(">>>>>>>>>>>>>>>>>>>>>>>>>>>>to build node_groups\n");
    }
    int current_flag = 0;
    int head = 0;
    MPI_Group world_group;
    MPI_Comm_group(MPI_COMM_WORLD, &world_group);
    while (group_split[head]>0){
        MPI_Group_incl(world_group, group_split[head], group_split+(head+1), node_groups+current_flag);
        MPI_Comm_create(MPI_COMM_WORLD,node_groups[current_flag],node_comms+current_flag);
        MPI_Comm head_comm;
        current_flag++;
        head = head+group_split[head]+1;
    }
    int inter_rank;
    MPI_Comm_rank(node_comms[my_flag],&inter_rank);
    // printf("global_rank=%d\tmy_flag=%d\tinter_rank=%d\n",rank,my_flag,inter_rank);
    int all_inter_ranks[MAX_NUM_NODE];
    if (rank==0){
        printf(">>>>>>>>>>>>>>>>>>>>>>>>>>>>to build head_group\n");
    }
    MPI_Allgather(&inter_rank,1,MPI_INT,all_inter_ranks,1,MPI_INT,MPI_COMM_WORLD);
    int head_groups_size = 0;
    for (int i = 0; i < nproc; i++){
        if (all_inter_ranks[i]==0){
            head_group_incl[head_groups_size++] = i;
        }
    }
    if (rank==0){
        printf("head_groups:\t");
        for (int i = 0; i < head_groups_size; i++){
            printf("%d\t", head_group_incl[i]);
        }
        printf("\n");
        
    }
    MPI_Group_incl(world_group, head_groups_size, head_group_incl, &head_group);
    MPI_Comm_create(MPI_COMM_WORLD, head_group, &head_comm);
    int head_rank=-1;
    if (inter_rank==0){
        MPI_Comm_rank(head_comm, &head_rank);
        printf("hostname=%s\tglobal_rank=%d\tinter_rank=%d\thead_rank=%d\n",
            hostname,rank,inter_rank,head_rank);
    }

    int data = 0;
    if (inter_rank==0 && head_rank==0){
        data = 9999;
    }
    MPI_Barrier(MPI_COMM_WORLD);
    printf("before bcast: rank=%d\thostname=%s\tdata=%d\n",rank,hostname,data);
    if (inter_rank==0){
        MPI_Bcast(&data,1,MPI_INT,0,head_comm);
    }
    MPI_Barrier(MPI_COMM_WORLD);
    printf("after bcast in head_group: rank=%d\thostname=%s\tdata=%d\n",rank,hostname,data);
    MPI_Bcast(&data,1,MPI_INT,0,node_comms[my_flag]);
    MPI_Barrier(MPI_COMM_WORLD);
    printf("after bcast in node_group: rank=%d\thostname=%s\tdata=%d\n",rank,hostname,data);
    MPI_Finalize();
    return 0;
}